#!/usr/local/bin/python

import argparse
import os
import sys
import hashlib
import socket
import psycopg2
import inspect

thisscript= inspect.getfile(inspect.currentframe())

parser = argparse.ArgumentParser(description= """

DESCRIPTION:
                                 
    Insert into table fastqfiles the files passed as arg. Other columns in table
    fastqfiles are filled in by interactive input.    

EXAMPLES
    
    insert_sblab_fastqfiles.py -f file1.fq.gz file2.fq.gz
    
    ## Use for-loop to read list of files from ls or similar command:
    > for f in `ls *.fq.gz`
    > do
    > insert_sblab_fastqfiles.py -f $f
    > done


TODO:
   
   - Input from pipe/stdin does not work (E.g. ls *bam | insert_sblab_fastqfiles.py -f -)
     because of raw_input() see also http://www.nerdparadise.com/tech/python/eofexceptionrawinput/
   
""", formatter_class= argparse.RawTextHelpFormatter)

parser.add_argument('-f', '--files',
                   type= str,
                   nargs= '+',
                   required= True,
                   help= '''File(s) to send to postgres.
                   
                   ''')

parser.add_argument('--connection', '-c',
                   type= str,
                   required= False,
                   help= '''A string to be passed to psycopg2.connect() to connect to the database. E.g. "dbname=testdb user=me password= mypwd host= localhost"
If None, connection parameters will be read from ~/.psycopgpass

''')

args = parser.parse_args()

# --------------------------[ Functions ]---------------------------------------

def sumfile(fobj):
    '''Returns an md5 hash for an object with read() method.'''
    m = hashlib.md5()
    while True:
        d = fobj.read(8096)
        if not d:
            break
        m.update(d)
    return m.hexdigest()


def md5sum(fname):
    '''Returns an md5 hash for file fname, or stdin if fname is "-".'''
    if fname == '-':
        ret = sumfile(sys.stdin)
    else:
        try:
            f = open(fname, 'rb')
        except:
            return 'Failed to open file %s' %fname
        ret = sumfile(f)
        f.close()
    return ret
# ------------------------------------------------------------------------------

if args.connection is None:
    try:
        conn= open(os.path.join(os.getenv("HOME"), '.psycopgpass'))
    except:
        sys.exit('I cannot read file %s and -c/--connection parameter is None.\nI cannot connect to any database.' %(os.path.join(os.getenv("HOME"), '.psycopgpass')))
    conn= conn.readlines()
    conn_args= [x.strip() for x in conn if x.strip() != '' and x.strip().startswith('#') is False][0]
    print('\nNOTICE: Connecting to database using parameters in ~/.psycopgpass\n')
else:
    conn_args= args.connection

conn = psycopg2.connect(conn_args)
cur= conn.cursor()

## Use for-loop on sh instead
#if args.files == ['-']:
#    args.files= sys.stdin.readlines()
#    args.files= [x.strip() for x in args.files]
# args.files= sorted(list(set(args.files)))

## Check all the files exist
for f in args.files:
    if not os.path.isfile(f):
       sys.exit('\n%s: File %s not found\n' %(thisscript, f)) 


for f in args.files:
    cur.execute('select count(*) from fastqfiles where fastqfile = %s', (f,))
    nrows= cur.fetchall()[0][0]
    if nrows > 0:
        print('Warning: %s found in table. Skipping.' %(f,))
        continue
    print('\n---------------[ Input for: %s ]-------------\n' %(f,))
    while True:
        ## This list will be passed to target table:
        record= [f]
        rowin= raw_input('encoding [<Sanger>, Illumina 1.0, Illumina 1.3+, Illumina 1.5+, FASTA]: ')
        if rowin == '':
            rowin= 'Sanger'
        record.append(rowin)
        
        record.append(md5sum(f))
        
        rowin= raw_input('description [None]: ')
        if rowin == '':
            rowin= None
        record.append(rowin)
        
        rowin= raw_input('service_id [None]: ')
        if rowin == '':
            rowin= None
        record.append(rowin)
    
        rowin= raw_input('barcode_id_deprecated [None]: ')
        if rowin == '':
            rowin= None
        record.append(rowin)
    
        rowin= raw_input('library_id [None]: ')
        if rowin == '':
            rowin= None
        record.append(rowin)
    
        insert_str= cur.mogrify("""INSERT INTO fastqfiles VALUES (%s, %s, %s, %s, %s, %s, %s)""", record )
        print('\n' + insert_str)
        try:
            cur.execute(insert_str)
        except Exception, e:
            print('\nTry to correct this error:')
            print(e)
            conn.commit()
            continue
        conn.commit()
        print('Record saved')
        break
cur.close()
conn.close()

sys.exit()